#include "service/client_api.h"

#include "service/service.h"
#include "smart_serve_log.h"

namespace OHOS {
namespace SmartServe {

std::string MakeErrorResponse(ErrorType type,
                              std::optional<std::string> customMessage = std::nullopt,
                              std::optional<std::string> param = std::nullopt)
{
    auto it = errorTable.find(type);
    if (it == errorTable.end()) {
        return json(ErrorResponse{.error = ErrorResponse::ErrorInfo{.type = "unknown_error",
                                                                    .message = "Unknown error",
                                                                    .param = param.value_or("null")}})
            .dump();
    }
    const auto& info = it->second;
    return json(ErrorResponse{.error = ErrorResponse::ErrorInfo{.type = info.first,
                                                                .message = customMessage.value_or(info.second),
                                                                .param = param.value_or("null")}})
        .dump();
}

std::string ClientAPIHandler::HandleRequest(const std::string& requestStr)
{
    SMART_SERVE_LOGI("[HandleRequest] %s", requestStr.c_str());

    Request request;
    try {
        auto requestJson = json::parse(requestStr);
        request = requestJson.get<Request>();
    } catch (const json::exception& e) {
        SMART_SERVE_LOGE("[ClientAPI] Json parse error: %s", e.what());
        return MakeErrorResponse(ErrorType::BAD_REQUEST_ERROR, "The request cannot be parsed.");
    }

    if (request.type == Request::EVAL_LLM) {
        return HandleEvalLLMRequest(request);
    }

    if (request.type == Request::GET_LLM_RESULT) {
        return HandleGetLLMResultRequest(request);
    }

    if (request.type == Request::CANCEL) {
        return HandleCancelRequest(request);
    }

    if (request.type == Request::LIST_MODELS) {
        return HandleListModelsRequest();
    }

    return MakeErrorResponse(ErrorType::BAD_REQUEST_ERROR, "The request type is not supported.");
}

// ========= EvalLLM =========

std::string ClientAPIHandler::HandleEvalLLMRequest(const Request& request)
{
    EvalLLMRequest evalRequest;
    try {
        evalRequest = request.request.get<EvalLLMRequest>();
    } catch (const json::exception& e) {
        SMART_SERVE_LOGE("[HandleEvalLLMRequest] Json parse error: %s", e.what());
        return MakeErrorResponse(ErrorType::BAD_REQUEST_ERROR, "The request cannot be parsed");
    }

    auto task = std::make_shared<EvalLLMTask>(request.pid, evalRequest.model, evalRequest.input);

    if (!EvalLLM(task)) {
        return MakeErrorResponse(ErrorType::NOT_FOUND_ERROR);
    }

    auto tid = tasks_.Add(std::move(task));

    EvalLLMResponse evalResponse{.id = tid};
    return json(evalResponse).dump();
}

bool ClientAPIHandler::EvalLLM(std::shared_ptr<EvalLLMTask> task)
{
    auto service = Service::Instance();

    // TODO: handle pid
    auto session = service->sessionManager_.GetOrCreateSession(task.get());
    if (!session) {
        return false; // ErrorType::NOT_FOUND_ERROR
    }
    session->AddTask(std::move(task));

    service->scheduler_->Schedule(std::make_unique<ScheduleEvent>());
    return true;
}

// ========= EvalLLM =========

// ========= GetLLMResult =========

std::string ClientAPIHandler::HandleGetLLMResultRequest(const Request& request)
{
    GetLLMResultRequest getRequest;
    try {
        getRequest = request.request.get<GetLLMResultRequest>();
    } catch (const json::exception& e) {
        SMART_SERVE_LOGE("[HandleGetLLMResultRequest] Json parse error: %s", e.what());
        return MakeErrorResponse(ErrorType::BAD_REQUEST_ERROR, "The request cannot be parsed");
    }

    auto baseTask = tasks_.Get(getRequest.id);
    if (!baseTask) {
        return MakeErrorResponse(ErrorType::NOT_FOUND_ERROR);
    }

    auto task = std::static_pointer_cast<EvalLLMTask>(baseTask);

    std::string content;
    {
        std::unique_lock<std::mutex> lock(task->mutex_);
        task->WaitForRemainingResult(lock);
        content = task->GetRemainingResult();
    }

    std::string error = task->doneReason();

    // 从TaskCenter删除并回收task
    if (task->isDone()) {
        tasks_.Remove(getRequest.id);
    }

    GetLLMResultResponse response{
        .id = getRequest.id,
        .content = content,
        .error = error,
    };

    std::string ret;
    try {
        ret = json(response).dump();
    } catch (const json::exception& e) {
        SMART_SERVE_LOGE("[HandleGetLLMResultRequest] response dump error: %s", e.what());
        return MakeErrorResponse(ErrorType::BAD_REQUEST_ERROR, "The request cannot be parsed");
    }

    return ret;
}

// ========= GetLLMResult =========

// ========= Cancel =========
std::string ClientAPIHandler::HandleCancelRequest(const Request& request)
{
    CancelRequest cancelRequest;
    try {
        cancelRequest = request.request.get<CancelRequest>();
    } catch (const json::exception& e) {
        SMART_SERVE_LOGE("[HandleCancelRequest] Json parse error: %s", e.what());
        return MakeErrorResponse(ErrorType::BAD_REQUEST_ERROR, "The request cannot be parsed");
    }

    auto task = tasks_.Get(cancelRequest.id);
    if (!task) {
        return MakeErrorResponse(ErrorType::NOT_FOUND_ERROR);
    }

    {
        std::unique_lock<std::mutex> lock(task->mutex_);
        task->state_ = EvalTask::TaskState::CANCELLED;
    }
    tasks_.Remove(cancelRequest.id);

    CancelResponse response{
        .status = "success",
        .id = cancelRequest.id,
    };
    return json(response).dump();
}
// ========= Cancel =========

// ========= ListModels =========
std::string ClientAPIHandler::HandleListModelsRequest()
{
    auto service = Service::Instance();
    ListModelsResponse response;
    response.object = "list";
    {
        std::unique_lock<std::mutex> lock(service->msMgmtMutex_);
        const auto& registry = service->sessionManager_.GetRegistry();
        for (const auto& [key, sessionCreator] : registry) {
            SessionInfo info;
            info.id = sessionCreator->Name();
            response.data.push_back(std::move(info));
        }
    }
    return json(response).dump();
}
// ========= ListModels =========

} // namespace SmartServe
} // namespace OHOS
